// Copyright 2019 JD.com Inc. JD AI

#include <dabnn/im2col.h>

#include <gtest/gtest.h>

#include <common/helper.h>
#include <common/log_helper.h>
#include <dabnn/mat.h>

namespace bnn {

TEST(im2col, im2col) {
    const size_t len = 32;
    float data[len];
    for (size_t i = 0; i < len; i++) {
        data[i] = i;
    }
    Mat im(4, 4, 2, data, DataType::Float);
    Mat col(288, DataType::Float);
    im2col(im, 3, 3, 1, 1, 1, 1, 1, 1, col);
    float data_col_expected[]{
        0,  0,  0,  0,  0,  0,  0,  0,  0,  1,  2,  3,  0,  0,  8,  9,  10, 11,
        0,  0,  0,  0,  0,  0,  0,  1,  2,  3,  4,  5,  8,  9,  10, 11, 12, 13,
        0,  0,  0,  0,  0,  0,  2,  3,  4,  5,  6,  7,  10, 11, 12, 13, 14, 15,
        0,  0,  0,  0,  0,  0,  4,  5,  6,  7,  0,  0,  12, 13, 14, 15, 0,  0,
        0,  0,  0,  1,  2,  3,  0,  0,  8,  9,  10, 11, 0,  0,  16, 17, 18, 19,
        0,  1,  2,  3,  4,  5,  8,  9,  10, 11, 12, 13, 16, 17, 18, 19, 20, 21,
        2,  3,  4,  5,  6,  7,  10, 11, 12, 13, 14, 15, 18, 19, 20, 21, 22, 23,
        4,  5,  6,  7,  0,  0,  12, 13, 14, 15, 0,  0,  20, 21, 22, 23, 0,  0,
        0,  0,  8,  9,  10, 11, 0,  0,  16, 17, 18, 19, 0,  0,  24, 25, 26, 27,
        8,  9,  10, 11, 12, 13, 16, 17, 18, 19, 20, 21, 24, 25, 26, 27, 28, 29,
        10, 11, 12, 13, 14, 15, 18, 19, 20, 21, 22, 23, 26, 27, 28, 29, 30, 31,
        12, 13, 14, 15, 0,  0,  20, 21, 22, 23, 0,  0,  28, 29, 30, 31, 0,  0,
        0,  0,  16, 17, 18, 19, 0,  0,  24, 25, 26, 27, 0,  0,  0,  0,  0,  0,
        16, 17, 18, 19, 20, 21, 24, 25, 26, 27, 28, 29, 0,  0,  0,  0,  0,  0,
        18, 19, 20, 21, 22, 23, 26, 27, 28, 29, 30, 31, 0,  0,  0,  0,  0,  0,
        20, 21, 22, 23, 0,  0,  28, 29, 30, 31, 0,  0,  0,  0,  0,  0,  0,  0};

    Mat col_expected(288, data_col_expected, DataType::Float);
    ASSERT_EQ(col.flatten(), col_expected.flatten());
}

}  // namespace bnn
